#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2; 


template<typename TYPE_Y, typename TYPE_DY, typename TYPE_Z> class KernelAsinhGrad {
    using T = TYPE_Y;
public:
    __aicore__ inline KernelAsinhGrad() {}
    __aicore__ inline void Init(GM_ADDR y, GM_ADDR dy, GM_ADDR z,
                                int32_t CoreDataNum, int32_t finalTileNum, uint64_t tileDataNum, int32_t TailDataNum,TPipe * pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->tileDataNum = tileDataNum;
        this->coreDataNum = CoreDataNum;
        this->tileNum = finalTileNum;
        this->tailDataNum = TailDataNum;

        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, this->coreDataNum);
        dyGm.SetGlobalBuffer((__gm__ DTYPE_DY*)dy, this->coreDataNum);
        zGm.SetGlobalBuffer((__gm__ DTYPE_Z*)z, this->coreDataNum);

        // pipe alloc memory to queue, the unit is Bytes



        pipe = pipeIn;

        pipe->InitBuffer(inQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));
        pipe->InitBuffer(inQueueDY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_DY));
        pipe->InitBuffer(outQueueZ, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Z));
        
        
        //pipe.InitBuffer(tmpBuffer1, this->tileDataNum * sizeof(DTYPE_Y));

        if constexpr (std::is_same_v<T, half>)
        {
            pipe->InitBuffer(tmp1,this->tileDataNum * sizeof(float));
            pipe->InitBuffer(tmp2,this->tileDataNum * sizeof(float));
            pipe->InitBuffer(tmpBuffer0, this->tileDataNum * sizeof(float));
            
            //pipe.InitBuffer(tmp3, this->tileDataNum * sizeof(float));
            //pipe.InitBuffer(tmp4, this->tileDataNum * sizeof(float));
        }
    }
    __aicore__ inline void Process() {
        int32_t loopCount = this->tileNum;
        this->processDataNum = this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            if (i == this->tileNum - 1) {
            this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<DTYPE_Y> yLocal = inQueueY.AllocTensor<DTYPE_Y>();
        LocalTensor<DTYPE_DY> dyLocal = inQueueDY.AllocTensor<DTYPE_DY>();

        DataCopy(yLocal, yGm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(dyLocal, dyGm[progress * this->tileDataNum], this->processDataNum);

        inQueueY.EnQue(yLocal);
        inQueueDY.EnQue(dyLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
        LocalTensor<DTYPE_Y> yLocal = inQueueY.DeQue<DTYPE_Y>();
        LocalTensor<DTYPE_DY> dyLocal = inQueueDY.DeQue<DTYPE_DY>();
        LocalTensor<DTYPE_Z> zLocal = outQueueZ.AllocTensor<DTYPE_Z>();

        //LocalTensor<float> tmpB0 = tmpBuffer0.Get<float>();
       
        //LocalTensor<DTYPE_Y> tmpB1 = tmpBuffer1.Get<DTYPE_Y>();

        if constexpr (std::is_same_v<DTYPE_Y, half>)
        {
            auto p1 = tmp1.Get<float>();
            auto p2 = tmp2.Get<float>();

            auto tmpB0 = tmpBuffer0.Get<float>();
            
            //auto p3 = tmp3.Get<float>();
            //auto p4 = tmp4.Get<float>();

            Cast(p1, yLocal, RoundMode::CAST_NONE, this->processDataNum);
            Cast(p2,dyLocal, RoundMode::CAST_NONE, this->processDataNum);

            Exp(p1,p1,this->processDataNum);
            Mul(tmpB0,p1,p1,this->processDataNum);

            Adds(tmpB0,tmpB0,static_cast<float>(1.0),this->processDataNum);

            Add(p1,p1,p1,this->processDataNum);

            Mul(p1,p1,p2,this->processDataNum);

            Div(p1,p1,tmpB0,this->processDataNum);

            Cast(zLocal, p1, RoundMode::CAST_NONE, this->processDataNum);
            



            



            





            
        }
        else
        {
            // if(progress !=this->tileNum - 1){

                Exp(yLocal,yLocal,this->processDataNum);

                Mul(zLocal,yLocal,yLocal,this->processDataNum);
                Adds(zLocal,zLocal,static_cast<DTYPE_Y>(1.0),this->processDataNum);
                Add(yLocal,yLocal,yLocal,this->processDataNum);
                Mul(yLocal,yLocal,dyLocal,this->processDataNum); //就这行代码有问题
                Div(zLocal,yLocal,zLocal,this->processDataNum);
            //  }

            // else{

            //     uint64_t mask = 8;
            //     const uint8_t number = this->tailDataNum /mask;

            //     Exp(yLocal,yLocal,this->tailDataNum);

            //     Mul(zLocal,yLocal,yLocal,this->tailDataNum);
            //     Adds(zLocal,zLocal,static_cast<DTYPE_Y>(1.0),this->tailDataNum);
            //     Add(yLocal,yLocal,yLocal,this->tailDataNum);
            //     Mul(yLocal,dyLocal,yLocal,mask,number,{1,1,1,8,8,8}); //就这行代码有问题
            //     Div(zLocal,yLocal,zLocal,this->tailDataNum);

                

                // Exp(tmpB0,yLocal,this->processDataNum);
                // Muls(zLocal,yLocal,static_cast<DTYPE_Y>(-1.0),this->processDataNum);
                // Exp(zLocal,zLocal,this->processDataNum);
                // Add(zLocal,tmpB0,zLocal,this->processDataNum);
                // Muls(dyLocal,dyLocal,static_cast<DTYPE_Y>(2.0),this->processDataNum);
                // Div(zLocal,dyLocal,zLocal,this->processDataNum);

                
                
            //}
        }
        
        outQueueZ.EnQue<DTYPE_Z>(zLocal);
        inQueueY.FreeTensor(yLocal);
        inQueueDY.FreeTensor(dyLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<DTYPE_Z> zLocal = outQueueZ.DeQue<DTYPE_Z>();
        DataCopy(zGm[progress * this->tileDataNum], zLocal, this->processDataNum);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    TPipe * pipe;
    // TBuf<QuePosition::VECCALC> tmpBuffer, signbitBuffer;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueY, inQueueDY;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
    TBuf<QuePosition::VECCALC> tmp1, tmp2, tmpBuffer0;

    GlobalTensor<DTYPE_Y> yGm;
    GlobalTensor<DTYPE_DY> dyGm;
    GlobalTensor<DTYPE_Z> zGm;
    int32_t coreDataNum;
    int32_t tileNum;
    uint64_t tileDataNum;
    int32_t tailDataNum;
    int32_t processDataNum;
};

extern "C" __global__ __aicore__ void asinh_grad(GM_ADDR y, GM_ADDR dy, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);

    TPipe pipe;

    KernelAsinhGrad<DTYPE_Y, DTYPE_DY, DTYPE_Z> op;
    op.Init(y, dy, z, tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum,&pipe);  
    op.Process();
    // TODO: user kernel impl
}